package cn.yhp.chatgpt.data.domain.openai.service;

import cn.bugstack.chatglm.model.ChatCompletionRequest;
import cn.bugstack.chatglm.model.ChatCompletionResponse;
import cn.bugstack.chatglm.model.EventType;
import cn.bugstack.chatglm.model.Model;
import cn.yhp.chatgpt.data.domain.openai.model.aggregates.ChatProcessAggregate;
import cn.yhp.chatgpt.data.domain.openai.model.entity.RuleLogicEntity;
import cn.yhp.chatgpt.data.domain.openai.model.valobj.LogicCheckTypeVO;
import cn.yhp.chatgpt.data.domain.openai.service.rule.ILogicFilter;
import cn.yhp.chatgpt.data.domain.openai.service.rule.factory.DefaultLogicFactory;
import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.core.JsonProcessingException;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author yhp
 * @description
 * @date 2024/2/2
 */
@Service
@Slf4j
public class ChatService extends AbstractChatService{

    @Resource
    private DefaultLogicFactory logicFactory;

    @Override
    protected void doMessageResponse(ChatProcessAggregate chatProcess, ResponseBodyEmitter emitter) throws JsonProcessingException {
        // 1、构建请求的问题消息
        List<ChatCompletionRequest.Prompt> prompts = chatProcess.getMessages().stream().map(e ->
                ChatCompletionRequest.Prompt.builder()
                        .role(e.getRole())
                        .content(e.getContent()).build()
        ).collect(Collectors.toList());
        // 2、封装请求,注意这里不用builder,有些参数需初始化值
        ChatCompletionRequest chatCompletionRequest=new ChatCompletionRequest();
        chatCompletionRequest
                // 暂时写死
                .setModel(Model.CHATGLM_TURBO);
        chatCompletionRequest.setPrompt(prompts);
        // 3、请求应答(调用ai接口)
        openAiSession.completions(chatCompletionRequest, new EventSourceListener() {

            @Override
            public void onClosed(@NotNull EventSource eventSource) {
                log.info("流式对话完成...");
            }

            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
                ChatCompletionResponse completionResponse = JSON.parseObject(data, ChatCompletionResponse.class);
                // log.info("数据：{}",completionResponse.getData());
                // 消息类型type ,add 增量，finish 结束，error 错误，interrupted 中断
                if (EventType.finish.getCode().equals(type)){
                    emitter.complete();
                    log.info("输出结束 tokens:{}",JSON.toJSONString(completionResponse.getMeta()));
                }
                try {
                    emitter.send(completionResponse.getData());
                    if (EventType.error.getCode().equals(type)){
                        emitter.complete();
                        log.info("输出异常 :{}",JSON.toJSONString(completionResponse.getData()));
                    }
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }

    @Override
    protected RuleLogicEntity<ChatProcessAggregate> doCheckLogic(ChatProcessAggregate chatProcess, DefaultLogicFactory.LogicModel... logics) throws Exception {
        Map<DefaultLogicFactory.LogicModel, ILogicFilter> logicFilterMap = logicFactory.openLogicFilter();
        RuleLogicEntity<ChatProcessAggregate> entity=null;
        // 将参数进行每个规则过滤
        for (DefaultLogicFactory.LogicModel logic : logics) {
            entity=logicFilterMap.get(logic).filter(chatProcess);
            // 不通过直接返回
            if (!LogicCheckTypeVO.SUCCESS.equals(entity.getType())){
                return entity;
            }
        }
        // 没有规则直接返回成功
        return entity!=null ? entity:RuleLogicEntity.<ChatProcessAggregate>builder()
                .type(LogicCheckTypeVO.SUCCESS).data(chatProcess).build();
    }
}
